#include "mdfu.h"
#include "uart_transport.h"
#include "mdfu_errors.h"
#include "ClientInfo.h"
#include <iostream>
#include <set>
#include <vector>
#include <cstddef>
#include <sstream>
#include <iomanip>
#include <sstream>

Mdfu::Mdfu() : retries(3), device_port("/dev/ttyUSB0")
{
    // 可以在此初始化更多的成员变量
    sequence_number = 0;
    std::cout << "Mdfu object created, device port set to " << device_port << std::endl;
}

Mdfu::Mdfu(int init_retries, std::string init_device_port)
    : retries(init_retries), device_port(init_device_port)
{
    sequence_number = 0;
    std::cout << "Mdfu initialized with retries: " << retries << " and device port: " << device_port << std::endl;
}

Mdfu::~Mdfu()
{
    std::cout << "Mdfu object destroyed." << std::endl;
}

void Mdfu::run_upgrade(const std::vector<uint8_t> &image)
{
    std::cout << "Running firmware upgrade with image of size " << image.size() << " bytes." << std::endl;

    std::cout << std::endl;
    int client_buffer_size = 5; // 假设 client_buffer_size 是一个固定值
    auto chunks = chunkify(image, client_buffer_size);

    // TODO 根据client info判断version、timeout、buffer size等参数
    get_client_info();

    start_transfer();
    for (const auto &chunk : chunks)
    {
        write_chunk(chunk);
    }
    uint8_t image_state = get_image_state();
    if (image_state != static_cast<uint8_t>(ImageState::VALID))
    {
        std::cout << "Image transfer failed with image state: " << static_cast<int>(image_state) << std::endl;
    }
    end_transfer();
}

void Mdfu::get_client_info()
{
    std::cout << "Getting client info." << std::endl;
    MdfuStatusPacket status_packet = send_cmd(MdfuCmd::GET_CLIENT_INFO); // 发送 GET_CLIENT_INFO 命令

    ClientInfo clientInfo = ClientInfo::fromBytes(status_packet.data); // 解析 ClientInfo 数据

    std::cout << "Parsed Client Info:" << std::endl;
    std::cout << "Protocol Version: " << clientInfo.getProtocolVersion() << std::endl;
    std::cout << "Buffer Count: " << static_cast<int>(clientInfo.getBufferCount()) << std::endl;
    std::cout << "Buffer Size: " << clientInfo.getBufferSize() << std::endl;
    std::cout << "Default Timeout: " << clientInfo.getDefaultTimeout() << " seconds" << std::endl;
    std::cout << "Inter Transaction Delay: " << clientInfo.getInterTransactionDelay() << " seconds" << std::endl;
}

void Mdfu::start_transfer(bool sync)
{
    std::cout << "Starting MDFU file transfer" << std::endl;
    send_cmd(MdfuCmd::START_TRANSFER, {}, sync); // 发送 START_TRANSFER 命令
}

void Mdfu::end_transfer()
{
    std::cout << "Ending MDFU file transfer" << std::endl;
    send_cmd(MdfuCmd::END_TRANSFER); // 发送 END_TRANSFER 命令
}

uint8_t Mdfu::get_image_state()
{
    std::cout << "Getting image state." << std::endl;
    MdfuStatusPacket status_packet = send_cmd(MdfuCmd::GET_IMAGE_STATE); // 发送 GET_IMAGE_STATE 命令

    if (status_packet.data.size() != 1) // 假设状态包至少有一个字节
    {
        std::cout << "Invalid status packet received." << std::endl;
        return 0;
    }

    return status_packet.data[0];
}

void Mdfu::write_chunk(const std::vector<uint8_t> &chunk)
{
    std::cout << "Writing chunk of size " << chunk.size() << " bytes to device." << std::endl;
    send_cmd(MdfuCmd::WRITE_CHUNK, chunk);
}

// TODO 目前是顺序分块，需要验证与pymdfu库是否一致
std::vector<std::vector<uint8_t>> Mdfu::chunkify(const std::vector<uint8_t> &data, size_t chunk_size, int padding)
{
    std::vector<std::vector<uint8_t>> chunks;
    size_t num_chunks = (data.size() + chunk_size - 1) / chunk_size;

    for (size_t i = 0; i < num_chunks; ++i)
    {
        size_t start_index = i * chunk_size;
        size_t end_index = std::min(start_index + chunk_size, data.size());
        chunks.emplace_back(data.begin() + start_index, data.begin() + end_index);
    }

    // Handle padding if the last chunk is not complete
    if (padding != -1)
    {
        size_t last_chunk_size = data.size() % chunk_size;
        if (last_chunk_size > 0)
        {
            size_t padding_count = chunk_size - last_chunk_size;
            std::vector<uint8_t> &last_chunk = chunks.back();
            last_chunk.insert(last_chunk.end(), padding_count, static_cast<uint8_t>(padding));
        }
    }

    return chunks;
}

MdfuStatusPacket Mdfu::send_cmd(MdfuCmd command, const std::vector<uint8_t> &data, bool sync)
{

    // 打印待发送命令和数据
    std::cout << "send command " << static_cast<int>(command);
    if (sync)
        std::cout << " sync";
    std::cout << " data size: " << data.size() << " bytes." << std::endl;
    for (uint8_t c : data)
        std::cout << std::hex << std::setw(2) << std::setfill('0') << (static_cast<int>(c) & 0xFF) << " ";

    // 构造命令包 MdfuCmdPacket
    MdfuCmdPacket cmd_packet(sequence_number, command, data, sync);
    std::cout << "Sending MDFU command packet: " << std::endl;
    std::cout << cmd_packet.to_string() << std::endl;

    // 构造最大尝试次数，设备响应状态包
    int attempts = 1 + retries;
    MdfuStatusPacket status_packet;
    Uart_Transport transport;

    while (attempts > 0)
    {
        try
        {
            transport.write(cmd_packet.to_binary()); // 发送命令
            auto response = transport.read(10);      // 读取响应，假设timeout已经在transport中设置
            status_packet = MdfuStatusPacket::from_binary(response);

            std::cout << "Received MDFU status packet: " << std::endl;
            std::cout << status_packet.to_string() << std::endl;

            if (status_packet.resend)
            {
                std::cout << "Resending MDFU packet due to resend request." << std::endl;
                attempts--; // 重试次数减一
                continue;   // 如果需要重发，则继续下一次循环   
            }

            if (status_packet.status == static_cast<int>(MdfuStatus::SUCCESS))
            {
                sequence_number++; // 成功则递增序列号
                break;             // 成功接收响应，退出循环
            }

            std::cout << "Error in response. Status: " << status_packet.status << std::endl;
            sequence_number++; // 即使发生错误也递增序列号，准备下一个命令
            break;             // 如果返回状态不是成功，跳出循环
        }
        catch (const MdfuErrors::TransportError &e)
        {
            std::cout << "Exception during command transmission: " << e.what() << std::endl;
            attempts--; // 重试次数减一
        }
        catch (const std::exception &e)
        {
            std::cerr << "General exception during command transmission: " << e.what() << std::endl;
        }

        if (attempts == 0)
        {
            std::cout << "Failed to send command after " << retries + 1 << " attempts." << std::endl;
        }
    }
    return status_packet;
}

void Mdfu::_increment_sequence_number()
{
    sequence_number = (sequence_number + 1) & 0x1F;
}

std::string Mdfu::testJni() {
    return "test JNI";
}

std::string Mdfu::testFrameData() {
    std::vector<uint8_t> data; // 空数据
    MdfuCmdPacket cmd_packet(0, MdfuCmd::GET_CLIENT_INFO, data, true);

    std::vector<uint8_t> binary = cmd_packet.to_binary();
    std::vector<uint8_t> frame_data = Frame(binary).to_bytes();

    std::stringstream ss;
    ss << std::hex << std::uppercase << std::setfill('0');
    for (uint8_t byte : frame_data) {
        ss << std::setw(2) << static_cast<int>(byte);
    }
    std::string result = ss.str();
    return result;

}

std::vector<uint8_t> Mdfu::test_send_cmd(MdfuCmd command, const std::vector<uint8_t>& data , bool sync){
    // 构造命令包 MdfuCmdPacket
    MdfuCmdPacket cmd_packet(sequence_number, command, data, sync);
    _increment_sequence_number();
    std::cout << "Sending MDFU command packet: " << std::endl;
    std::cout << cmd_packet.to_string() << std::endl;

    // 构造最大尝试次数，设备响应状态包
    int attempts = 1 + retries;
    MdfuStatusPacket status_packet;
    Uart_Transport transport;

    std::vector<uint8_t> frame_data = Frame(cmd_packet.to_binary()).to_bytes();
    std::cout << "Frame data: ";
    for (uint8_t c : frame_data)
    {
        std::cout << std::hex << (c & 0xFF) << " ";
    }
    std::cout << std::endl;
    return frame_data;
}



// =============================== MdfuCmdPacket =============================================

MdfuCmdPacket::MdfuCmdPacket(int sequence_number, MdfuCmd command, const std::vector<uint8_t> &data, bool sync)
    : sync(sync), command(command), data(data), sequence_number(sequence_number)
{
    if (sequence_number > 31 || sequence_number < 0)
    {
        std::cout << "Invalid MDFU packet sequence number: " << sequence_number << std::endl;
        throw std::invalid_argument("Valid values for MDFU packet sequence number are 0...31");
    }

    // Assuming MdfuCmd enum class is already defined and has all necessary commands
    std::set<MdfuCmd> cmd_values = {MdfuCmd::GET_CLIENT_INFO, MdfuCmd::START_TRANSFER, MdfuCmd::WRITE_CHUNK, MdfuCmd::GET_IMAGE_STATE, MdfuCmd::END_TRANSFER};
    if (cmd_values.find(command) == cmd_values.end())
    {
        std::string err_msg = "Invalid MDFU command: " + std::to_string(static_cast<int>(command));
        std::cout << err_msg << std::endl;
    }
}

std::vector<uint8_t> MdfuCmdPacket::to_binary()
{
    std::vector<uint8_t> packet;
    // Calculate the sequence field with the sync bit
    uint8_t sequence_field = static_cast<uint8_t>(sequence_number | (sync ? 0x80 : 0x00));
    packet.push_back(sequence_field);

    // Add command to the packet
    packet.push_back(static_cast<uint8_t>(command));

    // Append data to the packet
    packet.insert(packet.end(), data.begin(), data.end());

    return packet;
}

std::string getCommandName(MdfuCmd cmd)
{
    switch (cmd)
    {
    case MdfuCmd::GET_CLIENT_INFO:
        return "GET_CLIENT_INFO";
    case MdfuCmd::START_TRANSFER:
        return "START_TRANSFER";
    case MdfuCmd::WRITE_CHUNK:
        return "WRITE_CHUNK";
    case MdfuCmd::GET_IMAGE_STATE:
        return "GET_IMAGE_STATE";
    case MdfuCmd::END_TRANSFER:
        return "END_TRANSFER";
    default:
        return "UNKNOWN_COMMAND";
    }
}

std::string MdfuCmdPacket::to_string() const
{
    std::stringstream ss;
    ss << "Command:         " << getCommandName(command) << " (" << std::hex << static_cast<int>(command) << ")\n"
       << "Sequence Number: " << sequence_number << "\n"
       << "Sync:            " << (sync ? "true" : "false") << "\n"
       << "Data:            ";
    // for (auto byte : data)
    // {
    //     ss << std::hex << std::setfill('0') << std::setw(2) << (static_cast<int>(byte) & 0xFF);
    // }
    if (data.empty())
    {
        ss << "b''";
    }
    else
    {
        ss << "0x";
        for (uint8_t byte : data)
        {
            ss << std::hex << std::setfill('0') << std::setw(2)
               << (static_cast<uint8_t>(byte) & 0xFF);
        }
    }
    return ss.str();
}

// =============================== MdfuStatusPacket =============================================       
    
// TODO 待验证：接收到的MdfuStatusPacket是否正确、解析是否正确？
MdfuStatusPacket::MdfuStatusPacket(int sequence_number, int status, const std::vector<uint8_t> &data, bool resend)
    : sequence_number(sequence_number), status(status), data(data), resend(resend)
{
    if (sequence_number < 0 || sequence_number > 31)
    {
        throw std::invalid_argument("Sequence number must be between 0 and 31.");
    }

    // Define a set of all valid MdfuStatus values
    std::set<int> validStatuses = {static_cast<int>(MdfuStatus::SUCCESS),
                                   static_cast<int>(MdfuStatus::COMMAND_NOT_SUPPORTED),
                                   static_cast<int>(MdfuStatus::NOT_AUTHORIZED),
                                   static_cast<int>(MdfuStatus::COMMAND_NOT_EXECUTED),
                                   static_cast<int>(MdfuStatus::ABORT_FILE_TRANSFER)};

    // Check if the provided status is valid
    if (validStatuses.find(status) == validStatuses.end())
    {
        std::string error_message = "Invalid MDFU status: " + std::to_string(status);
        throw std::invalid_argument(error_message);
    }
}

std::string MdfuStatusPacket::to_string() const
{
    std::stringstream ss;
    ss << "Sequence Number: " << sequence_number << "\n"
       << "Status: " << status << "\n"
       << "Resend: " << resend << "\n"
       << "Data: 0x";
    for (uint8_t byte : data)
    {
        ss << std::hex << std::setfill('0') << std::setw(2) << (static_cast<unsigned>(byte) & 0xFF);
    }
    return ss.str();
}

std::vector<uint8_t> MdfuStatusPacket::to_binary() const
{
    std::vector<uint8_t> packet;
    uint8_t sequence_field = static_cast<uint8_t>(sequence_number | (resend ? 0x40 : 0x00));
    packet.push_back(sequence_field);
    packet.push_back(static_cast<uint8_t>(status));
    packet.insert(packet.end(), data.begin(), data.end());
    return packet;
}

std::tuple<int, int, std::vector<uint8_t>, bool> MdfuStatusPacket::decode_packet(const std::vector<uint8_t> &packet)
{
    if (packet.size() < 2)
    {
        throw std::invalid_argument("Packet too short to decode.");
    }
    int sequence_field = static_cast<unsigned char>(packet[0]);
    int sequence_number = sequence_field & 0x1F;
    bool resend = sequence_field & 0x40;
    int status = static_cast<unsigned char>(packet[1]);
    std::vector<uint8_t> data(packet.begin() + 2, packet.end());
    return {sequence_number, status, data, resend};
}

MdfuStatusPacket MdfuStatusPacket::from_binary(const std::vector<uint8_t> &packet)
{
    auto [sequence_number, status, data, resend] = decode_packet(packet);
    return MdfuStatusPacket(sequence_number, status, data, resend);
}
